import os
import thread
import time

import wx
import wx.lib.flatnotebook as fnb
import wx.grid as gridlib
import networkx
import matplotlib
from matplotlib.figure import Figure
from matplotlib.backends.backend_wxagg import \
    FigureCanvasWxAgg as FigCanvas
import numpy as np

# import task ?

class MyPanel(wx.Panel):

    def __init__(self, parent, task, **kwargs):

        self.task = task
        wx.Panel.__init__(self, parent, **kwargs)


class MyNB(wx.Notebook):

    def __init__(self, parent, task, **kwargs):

        self.task = task
        wx.Notebook.__init__(self, parent, **kwargs)


class HugeTable(gridlib.PyGridTableBase):
    def __init__(self):
        gridlib.PyGridTableBase.__init__(self)

        self.odd=gridlib.GridCellAttr()
        self.odd.SetBackgroundColour("sky blue")
        self.even=gridlib.GridCellAttr()
        self.even.SetBackgroundColour("sea green")

    def GetAttr(self, row, col, kind):
        attr = [self.even, self.odd][row % 2]
        attr.IncRef()
        return attr

    # This is all it takes to make a custom data table to plug into a
    # wxGrid.  There are many more methods that can be overridden, but
    # the ones shown below are the required ones.  This table simply
    # provides strings containing the row and column values.

    def GetNumberRows(self):
        return 10000

    def GetNumberCols(self):
        return 10000

    def IsEmptyCell(self, row, col):
        return False

    def GetValue(self, row, col):
        return str( (row, col) )

    def SetValue(self, row, col, value):
        self.log.write('SetValue(%d, %d, "%s") ignored.\n' % (row, col, value))


class ProteinTable(HugeTable):
    def __init__(self, data):
        HugeTable.__init__(self)
        self.col_label = ['id', 'start', 'end', 'strand', 'chromosome', 'taxon id']
        self.data = data

    def GetNumberRows(self):
        return len(self.data)

    def GetNumberCols(self):
        return 6

    def IsEmptyCell(self, row, col):
        return False

    def GetColLabelValue(self, col):
        #label = ['id', 'start', 'end', 'strand', 'chromosome', 'strand']
        return self.col_label[col]

    def GetValue(self, row, col):
        row_data = self.data[row].split('\t')
        return row_data[col]


class OrthTable(HugeTable):
    def __init__(self, data):
        HugeTable.__init__(self)
        self.col_label = ['id1', 'id2']
        self.data = data

    def GetNumberRows(self):
        return len(self.data)

    def GetNumberCols(self):
        return 2

    def IsEmptyCell(self, row, col):
        return False

    def GetColLabelValue(self, col):
        #label = ['id', 'start', 'end', 'strand', 'chromosome', 'strand']
        return self.col_label[col]

    def GetValue(self, row, col):
        row_data = self.data[row].split('\t')
        return row_data[col]


class ResultTable(HugeTable):
    def __init__(self, data):
        HugeTable.__init__(self)
        self.col_label = ['id1', 'id2', 'p-value', 'method']
        self.data = data

    def GetNumberRows(self):
        return len(self.data)

    def GetNumberCols(self):
        return 4

    def IsEmptyCell(self, row, col):
        return False

    def GetColLabelValue(self, col):
        #label = ['id', 'start', 'end', 'strand', 'chromosome', 'strand']
        return self.col_label[col]

    def GetValue(self, row, col):
        row_data = self.data[row].split('\t')
        return row_data[col]


class MyGrid(gridlib.Grid):

    def __init__(self, parent, task, **kwargs):
        self.task = task
        gridlib.Grid.__init__(self, parent, **kwargs)


class MyFrame(wx.Frame):
    '''self-define Frame, contains a Task instance
    '''
    def __init__(self, parent, task, **kwargs):

        self.task = task
        wx.Frame.__init__(self, parent, **kwargs)


class GraphPanel(MyPanel):
    def __init__(self, parent, task, **kwargs):
        MyPanel.__init__(self, parent, task)
        self.sizer = wx.BoxSizer(wx.HORIZONTAL)
        self.init_plot()
        self.sizer.Add(self.canvas, 1, wx.EXPAND)
        self.SetSizer(self.sizer)
        self.G = networkx.Graph()

    def init_plot(self):
        self.dpi = 100
        self.fig = Figure((3.0, 3.0), dpi=self.dpi)
        self.canvas = FigCanvas(self, -1, self.fig)

    def draw_graph(self):
        try:
            self.axes.clear()
        except:
            pass
        self.axes = self.fig.add_subplot(111)
        self.G = self.task.method.raw_result
        tmp_edges = self.G.edges(data=True)
        tmp_set_edges = set(tmp_edges)
        p = list()
        '''
        for i,j,k in tmp_edges:
            p.append(k)
        p.sort()
        if self.task.limit < len(p):
            cutoff = p[self.task.limit]
            self.G = networkx.Graph()
            for a,b,c in tmp_edges:
                if c < cutoff:        
                    self.G.add_node(b)
                    self.G.add_node(a)
                    self.G.add_edge(a,b)
        '''
        new_edges = list()
        for i in xrange(0,self.task.limit):
            new_edges.append(tmp_set_edges.pop())
        #new_edges = tmp_edges[0:self.task.limit]
        self.G = networkx.Graph()
        for a,b,c in new_edges:       
            self.G.add_node(b)
            self.G.add_node(a)
            self.G.add_edge(a,b)
            
        networkx.draw(self.G, 
                      ax=self.axes, 
                      node_size=self.task.node_size,
                      font_size=self.task.font_size,
                      node_color=self.task.node_color,
                      layout=self.task.layout
                      )
        self.canvas.draw()
        
        

class DrawSettingPanel(MyPanel):
    def __init__(self, parent, task, **kwargs):
        MyPanel.__init__(self, parent, task, **kwargs)
        self.parent = parent
        
        self.sizer = wx.BoxSizer(wx.VERTICAL)        

        box = wx.BoxSizer(wx.HORIZONTAL)
        label = wx.StaticText(self, -1, '  Node Size')
        box.Add(label, 2, wx.EXPAND)

        self.node_sc = wx.SpinCtrl(self, -1, "")
        self.node_sc.SetRange(1,100)
        self.node_sc.SetValue(5)

        box.Add(self.node_sc, 3, wx.EXPAND|wx.ALIGN_RIGHT)
        self.sizer.Add(box, 0, wx.ALIGN_LEFT|wx.EXPAND)

        box = wx.BoxSizer(wx.HORIZONTAL)

        label = wx.StaticText(self, -1, '  Node Color')
        box.Add(label, 2, wx.EXPAND)

        self.node_color_ch = wx.Choice(self, -1, choices=['black', 'red', 'green','blue'])
        self.Bind(wx.EVT_CHOICE, self.OnNodeColorChange, self.node_color_ch)
        box.Add(self.node_color_ch, 3, wx.EXPAND|wx.ALIGN_RIGHT)

        self.sizer.Add(box, 0, wx.ALIGN_LEFT|wx.EXPAND)

        box = wx.BoxSizer(wx.HORIZONTAL)
        label = wx.StaticText(self, -1, '  Node Shape')
        box.Add(label, 2, wx.EXPAND)
        self.node_shape_ch = wx.Choice(self, -1, choices=list('so^>v<dph8'))
        self.Bind(wx.EVT_CHOICE, self.OnNodeShapeChange, self.node_shape_ch)
        box.Add(self.node_shape_ch, 3, wx.EXPAND)
        self.sizer.Add(box, 0, wx.ALIGN_LEFT|wx.EXPAND)

        box = wx.BoxSizer(wx.HORIZONTAL)
        label = wx.StaticText(self, -1, '  Font Size')
        box.Add(label, 2, wx.EXPAND)
        self.font_sc = wx.SpinCtrl(self, -1, "")
        self.font_sc.SetRange(1,20)
        self.font_sc.SetValue(5)

        box.Add(self.font_sc, 3, wx.EXPAND|wx.ALIGN_RIGHT)
        self.sizer.Add(box, 0, wx.ALIGN_LEFT|wx.EXPAND)
        
        box = wx.BoxSizer(wx.HORIZONTAL)
        label = wx.StaticText(self, -1, '  Layout')
        box.Add(label, 2, wx.EXPAND)
        self.layout_ch = wx.Choice(self, -1, choices = ['spring'])
        self.Bind(wx.EVT_CHOICE, self.OnLayoutChange, self.layout_ch)
        box.Add(self.layout_ch, 3, wx.EXPAND)
        self.sizer.Add(box, 0, wx.ALIGN_LEFT|wx.EXPAND)

        box = wx.BoxSizer(wx.HORIZONTAL)
        label = wx.StaticText(self, -1, '  Edge limit')
        box.Add(label, 2, wx.EXPAND)

        self.limit_sc = wx.SpinCtrl(self, -1, "")
        self.limit_sc.SetRange(1,300)
        self.limit_sc.SetValue(20)

        box.Add(self.limit_sc, 3, wx.EXPAND|wx.ALIGN_RIGHT)
        self.sizer.Add(box, 0, wx.ALIGN_LEFT|wx.EXPAND)
        
        button = wx.Button(self, -1, 'Update')
        self.sizer.Add(button, 0, wx.LEFT)
        self.Bind(wx.EVT_BUTTON, self.OnUpdate, button)

        self.SetSizer(self.sizer)

    def OnNodeColorChange(self, e):
        self.task.set_node_color(e.GetString())

    def OnNodeShapeChange(self, e):
        self.task.set_node_shape(e.GetString())

    def OnLayoutChange(self, e):
        pass

    def OnUpdate(self, e):
        self.task.set_node_size(self.node_sc.GetValue())
        self.task.set_font_size(self.font_sc.GetValue())
        self.task.set_edge_limit(self.limit_sc.GetValue())
        #self.parent.parent.graph_panel.draw_graph()


class GCNFrame(MyFrame):
    '''the main Frame, contains everything
    '''
    def __init__(self, parent, task, **kwargs):
        MyFrame.__init__(self, parent, task, title='GCNBuilder', size=(800, 600), **kwargs)
        self.create_menu_bar()
        self.layout_items()

    def create_menu_bar(self):
        self._menu_bar = wx.MenuBar(wx.MB_DOCKABLE)
        self._task_menu = wx.Menu()
        self._edit_menu = wx.Menu()
        self._help_menu = wx.Menu()

        item = wx.MenuItem(self._task_menu, -1, 'New')
        self._task_menu.AppendItem(item)
        self.Bind(wx.EVT_MENU, self.OnNew, item)

        item = wx.MenuItem(self._task_menu, -1, 'Run')
        self._task_menu.AppendItem(item)
        self.Bind(wx.EVT_MENU, self.OnRun, item)

        item = wx.MenuItem(self._task_menu, -1, 'Draw')
        self._task_menu.AppendItem(item)
        self.Bind(wx.EVT_MENU, self.OnDraw, item)
        
        item = wx.MenuItem(self._task_menu, -1, 'Exit')
        self._task_menu.AppendItem(item)

        self._menu_bar.Append(self._task_menu, 'Task')
        self._menu_bar.Append(self._edit_menu, 'Edit')
        self._menu_bar.Append(self._help_menu, 'Help')

        self.SetMenuBar(self._menu_bar)

    def layout_items(self):
        book_style = fnb.FNB_NODRAG | fnb.FNB_NO_X_BUTTON

        main_sizer = wx.BoxSizer(wx.VERTICAL)
        self.SetSizer(main_sizer)

        self.up_panel = wx.Panel(self, -1)

        up_sizer = wx.BoxSizer(wx.HORIZONTAL)
        self.up_panel.SetSizer(up_sizer)

        self.setting_book = SettingNB(self.up_panel, self.task)

        up_sizer.Add(self.setting_book, 2, wx.EXPAND)

        self.graph_panel = GraphPanel(self.up_panel, self.task)
        up_sizer.Add(self.graph_panel, 5, wx.EXPAND)

        main_sizer.Add(self.up_panel, 6, wx.EXPAND)

        spacer = wx.Panel(self, -1)
        spacer.SetBackgroundColour(wx.SystemSettings_GetColour(wx.SYS_COLOUR_3DFACE))
        main_sizer.Add(spacer, 0, wx.ALL | wx.EXPAND)

        self.data_book = fnb.FlatNotebook(self, style = book_style)
        main_sizer.Add(self.data_book, 3, wx.EXPAND)

        self.Freeze()

        self.protein_grid = ProteinGrid(self.data_book, self.task)
        self.orth_grid = OrthGrid(self.data_book, self.task)
        self.data_book.AddPage(self.protein_grid, 'Protein Data')
        self.data_book.AddPage(self.orth_grid, 'Ortholog Data',select=False)
        self.Thaw()


    def OnRun(self, e):
        self.task.calculate()
        #self.setting_book.SetFocus(2)
        #self.graph_panel.draw_graph()

    def OnDraw(self, e):
        self.graph_panel.draw_graph()
        
    def OnNew(self, e):
        self.create_menu_bar()
        self.layout_items()
        self.task.renew()


class InputSettingPanel(MyPanel):
    def __init__(self, parent, task, **kwargs):
        MyPanel.__init__(self, parent, task, **kwargs)

        sizer = wx.BoxSizer(wx.VERTICAL)

        label = wx.StaticText(self, -1, '  Choose Methods')
        sizer.Add(label, 0)
        self.method_ch = wx.Choice(self, -1, (100, 50), choices=task.method_list)
        self.Bind(wx.EVT_CHOICE, self.MethodChoice, self.method_ch)
        sizer.Add(self.method_ch, 0, wx.ALIGN_LEFT)

        spacer = wx.Panel(self, -1, (20, 20))
        sizer.Add(spacer, 0)

        label = wx.StaticText(self, -1, '  Protein File')
        sizer.Add(label, 0, wx.EXPAND)

        box = wx.BoxSizer(wx.HORIZONTAL)

        self.protein_text = wx.TextCtrl(self, -1, '')
        box.Add(self.protein_text, 5, wx.ALIGN_LEFT)

        button = wx.Button(self, -1, '...')
        self.Bind(wx.EVT_BUTTON, self.OnProteinFileOpen, button)
        box.Add(button, 1, wx.ALIGN_RIGHT)

        sizer.Add(box, 0, wx.EXPAND)

        label = wx.StaticText(self, -1, '  Ortholog File')
        sizer.Add(label, 0, wx.ALIGN_LEFT)

        box = wx.BoxSizer(wx.HORIZONTAL)

        self.orth_text = wx.TextCtrl(self, -1, '')
        box.Add(self.orth_text, 5, wx.ALIGN_LEFT)

        button = wx.Button(self, -1, '...')
        self.Bind(wx.EVT_BUTTON, self.OnOrthFileOpen, button)
        box.Add(button, 1, wx.ALIGN_RIGHT)

        sizer.Add(box, 0, wx.EXPAND)

        label = wx.StaticText(self, -1, '  Pick a data storage strategy')
        sizer.Add(label, 0, wx.ALIGN_LEFT)

        self.db_ch = wx.Choice(self, -1, (100, 50), choices=['memory'])
        self.Bind(wx.EVT_CHOICE, self.DBChoice, self.db_ch)
        sizer.Add(self.db_ch, 0, wx.ALIGN_LEFT)

        button = wx.Button(self, -1, 'Load Data')
        self.Bind(wx.EVT_BUTTON, self.OnLoadData, button)
        sizer.Add(button, 0)

        spacer = wx.Panel(self, -1, (20, 20))
        sizer.Add(spacer, 0)

        label = wx.StaticText(self, -1, '  Select Taxon (after data loaded)')
        sizer.Add(label, 0, wx.ALIGN_LEFT)

        self.taxon_ch = wx.Choice(self, -1, (100, 50), choices = [])
        self.taxon_ch.Disable()
        self.Bind(wx.EVT_CHOICE, self.TaxonChoice, self.taxon_ch)
        sizer.Add(self.taxon_ch, 0, wx.ALIGN_LEFT)


        spacer = wx.Panel(self, -1)
        sizer.Add(spacer, 1, wx.EXPAND)

        self.SetSizer(sizer)

    def OnProteinFileOpen(self, event):
        dlg = wx.FileDialog(
            self, message="Choose a protein information file",
            defaultDir=os.getcwd(),
            defaultFile="",
            style=wx.OPEN | wx.CHANGE_DIR
            )

        if dlg.ShowModal() == wx.ID_OK:
            # This returns a Python list of files that were selected.
            self.task.set_protein_file(dlg.GetPath())
            self.protein_text.SetValue(self.task.protein_file)

        dlg.Destroy()

    def OnOrthFileOpen(self, event):
        dlg = wx.FileDialog(
            self, message="Choose a ortholog pair file",
            defaultDir=os.getcwd(),
            defaultFile="",
            style=wx.OPEN | wx.CHANGE_DIR
            )

        if dlg.ShowModal() == wx.ID_OK:
            # This returns a Python list of files that were selected.
            self.task.set_orth_file(dlg.GetPath())
            self.orth_text.SetValue(self.task.orth_file)

        dlg.Destroy()

    def OnLoadData(self, e):
        input_nb = self.GetParent()
        up_panel = input_nb.GetParent()
        main_frame = up_panel.GetParent()

        # update the grids, showing protein info and ortholog pair
        protein_grid = main_frame.protein_grid
        total_protein = protein_grid.update_data(self.task.protein_file)
        orth_grid = main_frame.orth_grid
        total_orth = orth_grid.update_data(self.task.orth_file)

        # build the database
        '''
        max = 100
        dlg = wx.ProgressDialog("Data dumping progress",
                       "",
                       maximum = max,
                       parent=self,
                       style = wx.PD_APP_MODAL
                        | wx.PD_ELAPSED_TIME
                        | wx.PD_REMAINING_TIME
                        )
        '''
        #db_thread = thread.start_new_thread(self.task.init_db, ())
        self.task.init_db()
        '''
        while True:
            if self.task.db:
                break
            
        while self.task.status != 'database set':
            time.sleep(0.5)
            #print self.task.db.progress
            if self.task.db.progress == 'protein':
                label = 'Dumping protein data'
                count = self.task.db.count*1.0/(total_protein+total_orth*2)*100
            else:
                label = 'Dumping Ortholog data'
                count = (self.task.db.count*1.0+total_protein)/(total_protein+total_orth*2)*100
            
            if count >= max:
                count -= 0.1
                label = 'Almost there, please wait'
            dlg.Update(count, label)
            
        dlg.Destroy()
        '''
        
        # update the taxon choice
        taxon_set = self.task.db.get_taxons()
        self.taxon_ch.AppendItems(list(taxon_set))
        self.taxon_ch.Enable()
        self.taxon_ch.Update()
        

    def updata_grid(self, protein_grid=None, orth_grid=None):
        self.protein_grid = protein_grid
        self.orth_grid = orth_grid

    def DBChoice(self, e):
        #self.task.set_db_name(e.GetString())
        pass

    def TaxonChoice(self, e):
        self.task.set_target_taxon(e.GetString())

    def MethodChoice(self, e):
        self.task.set_method(e.GetString())


class SettingNB(MyNB):
    def __init__(self, parent, task, **kwargs):
        self.parent = parent
        MyNB.__init__(self, parent, task, **kwargs)
        self.input_setting_panel = InputSettingPanel(self, self.task)
        self.AddPage(self.input_setting_panel, 'Input Setting')
        self.draw_setting_panel = DrawSettingPanel(self, self.task)
        self.AddPage(self.draw_setting_panel, 'Graph Setting')


class ProteinGrid(MyGrid):
    def __init__(self, parent, task, filename='', **kwargs):
        MyGrid.__init__(self, parent, task, **kwargs)
        self.update_data(filename)

        # The second parameter means that the grid is to take ownership of the
        # table and will destroy it when done.  Otherwise you would need to keep
        # a reference to it and call it's Destroy method later.


        #self.Bind(gridlib.EVT_GRID_CELL_RIGHT_CLICK, self.OnRightDown)

    def get_data(self, filename):
        self.data = []
        if filename != '':
            for line in open(filename):
                tmp = line.strip().split()
                self.data.append(line.strip())
        self.table = ProteinTable(self.data)


    def update_data(self, filename):
        self.protein_file = filename
        self.get_data(self.protein_file)
        self.SetTable(self.table, True)
        self.Update()
        
        return len(self.data)


class OrthGrid(MyGrid):
    def __init__(self, parent, task, filename='', **kwargs):
        MyGrid.__init__(self, parent, task, **kwargs)
        self.update_data(filename)


        # The second parameter means that the grid is to take ownership of the
        # table and will destroy it when done.  Otherwise you would need to keep
        # a reference to it and call it's Destroy method later.


        #self.Bind(gridlib.EVT_GRID_CELL_RIGHT_CLICK, self.OnRightDown)

    def get_data(self, filename):
        self.data = []
        if filename != '':
            for line in open(filename):
                self.data.append(line.strip())
        self.table = OrthTable(self.data)


    def update_data(self, filename):
        self.protein_file = filename
        self.get_data(self.protein_file)
        self.SetTable(self.table, True)
        self.Update()
        
        return len(self.data)


class ResultGrid(MyGrid):
    def __init__(self, parent, task, filename='', **kwargs):
        MyGrid.__init__(self, parent, task, **kwargs)
        self.update_data(filename)


        # The second parameter means that the grid is to take ownership of the
        # table and will destroy it when done.  Otherwise you would need to keep
        # a reference to it and call it's Destroy method later.


        #self.Bind(gridlib.EVT_GRID_CELL_RIGHT_CLICK, self.OnRightDown)

    def get_data(self, filename):
        self.data = []
        if filename != '':
            for line in open(filename):
                tmp = line.strip().split()
                self.data.append(line.strip())
        self.table = ResultTable(self.data)

    def update_data(self, filename):
        self.protein_file = filename
        self.get_data(self.protein_file)
        self.SetTable(self.table, True)
        self.Update()






#END